import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'

import csv
import pickle

import numpy as np
import tensorflow as tf


if __name__ == "__main__":
    input_path = "/home/ubuntu/Ascend/model_inference/tf_source/inverse_inputs/tf.image.rgb_to_hsv/tf.image.rgb_to_hsv.pickle"
    compare_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/inverse/tf.image.rgb_to_hsv/20230321_183958/tf.image.rgb_to_hsv.csv"
    source_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/inverse/tf.image.rgb_to_hsv/20230321_183958/tf.image.rgb_to_hsv_source.pickle"
    follow_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/inverse/tf.image.rgb_to_hsv/20230321_183958/tf.image.rgb_to_hsv_follow.pickle"

    with open(input_path, "rb") as f:
        input_dict_list = pickle.load(f)

    source_results = []
    follow_results = []
    compare_results = []

    cnt = 0
    for input_dict in input_dict_list:
        x = np.array([])
        for key in input_dict.keys():
            val = input_dict[key]
            if type(val) is np.ndarray:
                x = val
        try:
            a = tf.image.rgb_to_hsv(**input_dict)
            print("OP 1 output:")
            print(a)
            source_results.append(a.numpy())
        except Exception as e:
            print("[OP 1] Exception occurred~:\n", e)
            source_results.append("Error")
            follow_results.append("No Follow")
            compare_results.append([cnt, "Source Error"])
            cnt += 1
            continue

        try:
            b = tf.image.hsv_to_rgb(a)
            print("OP 2 output:")
            print(b)
            b = b.numpy()
            follow_results.append(b)
            if type(b) is not np.ndarray:
                compare_results.append([cnt, "Not Tensor Inconsistent"])
            elif x.size != b.size:
                compare_results.append([cnt, "Inconsistent", x.shape, b.shape])
            elif np.allclose(x, b, atol=1.e-2, rtol=1.e-5, equal_nan=True):
                compare_results.append([cnt, "Consistent", x.shape, b.shape, x.dtype, b.dtype])
            else:
                compare_results.append([cnt, "Inconsistent", x.shape, b.shape, x.dtype, b.dtype])
            cnt += 1
        except Exception as e:
            print("[OP 2] Exception occurred~:\n", e)
            follow_results.append("Error")
            compare_results.append([cnt, "Follow Error"])
            cnt += 1

    with open(source_result_path, "wb+") as f:
        pickle.dump(source_results, f)

    with open(follow_result_path, "wb+") as f:
        pickle.dump(follow_results, f)

    with open(compare_result_path, "w+", newline="") as f:
        w = csv.writer(f)
        w.writerows(compare_results)
